/**
 * @author: dn-jinmin/dn-jinmin
 * @doc:
 */

package main

import (
	"context"
	"fmt"
	"github.com/tmc/langchaingo/callbacks"
	"github.com/tmc/langchaingo/chains"
	"github.com/tmc/langchaingo/llms/openai"
	"github.com/tmc/langchaingo/outputparser"
	"github.com/tmc/langchaingo/prompts"
)

func main() {
	url := "https://api.aiproxy.io/v1"
	apiKey := ""

	ctx := context.Background()

	llm, err := openai.New(openai.WithBaseURL(url), openai.WithToken(apiKey))
	if err != nil {
		panic(err)
	}

	// 定义要求大模型返回的数据格式
	output := outputparser.NewStructured([]outputparser.ResponseSchema{
		{
			Name:        "city",
			Description: "城市",
		}, {
			Name:        "temperature",
			Description: "气温",
		}, {
			Name:        "weather",
			Description: "天气",
		},
	})
	// output会根据指定的内容生成要求大模型根据指定格式返回的提示词，与基础提示词拼接就是最终提示词
	template := "请问今天{{.input}}天气怎么样" + output.GetFormatInstructions()
	fmt.Println("提示词：", template)
	// 这是组成用户问题的基础提示词，其中city则是可代替的内容
	prompt := prompts.NewPromptTemplate(template, []string{"input"})
	// callback是langchain中的事件，可以打印出其中交互的提示词
	callback := callbacks.LogHandler{}
	chain := chains.NewLLMChain(llm, prompt, chains.WithCallback(callback))
	// 运行
	res, err := chains.Call(ctx, chain, map[string]any{
		"input": "长沙",
	})
	if err != nil {
		panic(err)
	}

	fmt.Println("结果: ", res)
}
